from .module import PREPROCESS


@PREPROCESS.register_module
class Normalize:

    def __init__(self, mean, std, input_channel='rgb'):
        self.mean = mean
        self.std = std
        self.input_channel = input_channel

    def process(self, img, anno):
        anno = anno.copy()
        if self.input_channel == 'bgr':
            img = self.channel_transform(img)
        for i in range(3):
            img[:, :, i] = (img[:, :, i] - self.mean[i]) / self.std[i]
        return img, anno

    def channel_transform(self, img):
        lc = img[:, :, 2]
        img[:, :, 2] = img[:, :, 0]
        img[:, :, 0] = lc
        return img